package org.movee.jalgo.search;

/**
 * 红黑树特点：
 * 1) 每个节点或是红色，或是黑色。
 * 2) 根节点是黑色。
 * 3）所有叶子都是黑色（叶子是NIL节点）。
 * 4）每个红色节点必须有两个黑色的子节点。（从每个叶子到根的所有路径上不能有两个连续的红色节点。）
 * 5）从任一节点到其每个叶子的所有简单路径都包含相同数目的黑色节点。
 *
 * @param <K>
 */
public class RbTree<K extends Comparable<K>> {

    private static final boolean BLACK = true;
    private static final boolean RED   = false;

    // 哨兵节点
    private final Node<K> NIL;
    private Node<K> root;    // 根结点

    public RbTree() {
        NIL = new Node<>(null, BLACK, null, null, null);
        root = NIL;
    }

    private void processNode(Node<K> node) {
        System.out.print(node.key + " ");
    }

    private void preOrder(Node<K> tree) {
        if (tree != NIL) {
            processNode(tree);
            preOrder(tree.left);
            preOrder(tree.right);
        }
    }

    public void preOrder() {
        preOrder(root);
    }

    /*
     * 中序遍历"红黑树"
     */
    private void inOrder(Node<K> tree) {
        if (tree != NIL) {
            inOrder(tree.left);
            processNode(tree);
            inOrder(tree.right);
        }
    }

    public void inOrder() {
        inOrder(root);
    }


    private void postOrder(Node<K> tree) {
        if (tree != NIL) {
            postOrder(tree.left);
            postOrder(tree.right);
            processNode(tree);
        }
    }

    public void postOrder() {
        postOrder(root);
    }


    private Node<K> search(Node<K> x, K key) {
        if (x == NIL) {
            return NIL;
        }

        int cmp = key.compareTo(x.key);
        if (cmp < 0) {
            return search(x.left, key);
        } else if (cmp > 0) {
            return search(x.right, key);
        } else {
            return x;
        }
    }

    public Node<K> search(K key) {
        return search(root, key);
    }

    private Node<K> iterativeSearch(Node<K> x, K key) {
        while (x != NIL) {
            int cmp = key.compareTo(x.key);
            if (cmp < 0) {
                x = x.left;
            } else if (cmp > 0) {
                x = x.right;
            } else {
                return x;
            }
        }

        return x;
    }

    public Node<K> iterativeSearch(K key) {
        return iterativeSearch(root, key);
    }

    private Node<K> minimum(Node<K> tree) {
        if (tree == NIL) {
            return NIL;
        }

        while (tree.left != NIL) {
            tree = tree.left;
        }
        return tree;
    }

    public K minimum() {
        Node<K> p = minimum(root);
        return p == NIL ? null : p.key;
    }

    private Node<K> maximum(Node<K> tree) {
        if (tree == NIL) {
            return NIL;
        }

        while (tree.right != NIL) {
            tree = tree.right;
        }
        return tree;
    }

    public K maximum() {
        Node<K> p = maximum(root);
        return p == NIL ? null : p.key;
    }

    public Node<K> successor(Node<K> x) {
        // 如果x存在右孩子，则"x的后继结点"为 "以其右孩子为根的子树的最小结点"。
        if (x.right != NIL) {
            return minimum(x.right);
        }

        // 如果x没有右孩子。则x有以下两种可能：
        // (01) x是"一个左孩子"，则"x的后继结点"为 "它的父结点"。
        // (02) x是"一个右孩子"，则查找"x的最低的父结点，并且该父结点要具有左孩子"，找到的这个"最低的父结点"就是"x的后继结点"。
        Node<K> parent = x.parent;
        while ((parent != NIL) && (x == parent.right)) {
            x = parent;
            parent = parent.parent;
        }

        return parent;
    }

    public Node<K> predecessor(Node<K> x) {
        // 如果x存在左孩子，则"x的前驱结点"为 "以其左孩子为根的子树的最大结点"。
        if (x.left != NIL) {
            return maximum(x.left);
        }

        // 如果x没有左孩子。则x有以下两种可能：
        // (01) x是"一个右孩子"，则"x的前驱结点"为 "它的父结点"。
        // (01) x是"一个左孩子"，则查找"x的最低的父结点，并且该父结点要具有右孩子"，找到的这个"最低的父结点"就是"x的前驱结点"。
        Node<K> parent = x.parent;
        while ((parent != NIL) && (x == parent.left)) {
            x = parent;
            parent = parent.parent;
        }

        return parent;
    }

    private void leftRotate(Node<K> x) {
        // 设置x的右孩子为y
        Node<K> y = x.right;

        // 将 “y的左孩子” 设为 “x的右孩子”；
        // 如果y的左孩子非空，将 “x” 设为 “y的左孩子的父亲”
        x.right = y.left;
        if (y.left != NIL) {
            y.left.parent = x;
        }

        // 将 “x的父亲” 设为 “y的父亲”
        y.parent = x.parent;

        if (x.parent == NIL) {
            root = y;            // x的父节点是NIL，说明x是根节点，将根节点设置为y
        } else {
            if (x.parent.left == x) {
                x.parent.left = y;    // 如果 x是它父节点的左孩子，则将y设为“x的父节点的左孩子”
            } else {
                x.parent.right = y;    // 如果 x是它父节点的左孩子，则将y设为“x的父节点的左孩子”
            }
        }

        // 将 “x” 设为 “y的左孩子”
        y.left = x;
        // 将 “x的父节点” 设为 “y”
        x.parent = y;
    }

    private void rightRotate(Node<K> y) {
        // 设置x是当前节点的左孩子。
        Node<K> x = y.left;

        // 将 “x的右孩子” 设为 “y的左孩子”；
        // 如果"x的右孩子"不为空的话，将 “y” 设为 “x的右孩子的父亲”
        y.left = x.right;
        if (x.right != NIL) {
            x.right.parent = y;
        }

        // 将 “y的父亲” 设为 “x的父亲”
        x.parent = y.parent;

        if (y.parent == NIL) {
            root = x;            // y的父节点是NIL，说明y是根节点，将根节点设置为x
        } else {
            if (y == y.parent.right) {
                y.parent.right = x;    // 如果 y是它父节点的右孩子，则将x设为“y的父节点的右孩子”
            } else {
                y.parent.left = x;    // (y是它父节点的左孩子) 将x设为“x的父节点的左孩子”
            }
        }

        // 将 “y” 设为 “x的右孩子”
        x.right = y;

        // 将 “y的父节点” 设为 “x”
        y.parent = x;
    }

    private void insertFixUp(Node<K> node) {

        Node<K> parent;
        Node<K> gparent; // grand parent

        // 不管node是root节点还是非root点，node.parent都不会为空，因为root节点的parent是NIL节点
        // 插入root节点和root的孩子节点，都不会进入while循环，因为root的父节点和root节点颜色都是黑色
        // 下面操作完后，node节点都是红色，黑高都是平衡的，所以只要判断父节点是红色则进入循环
        // 每一次循环结束后，node节点都不会是NIL节点
        // 如果循环结束后node是root节点或root的孩子节点，不会再次进入while循环了，因为它的父节点是黑色的
        while (node.parent.color == RED) {
            // 进入循环的node，parent和gparent都不会是NIL，因为root和root的孩子节点不会进入循环
            parent = node.parent;
            gparent = parent.parent;

            //若“父节点”是“祖父节点的左孩子”
            if (parent == gparent.left) {
                // Case 1条件：叔叔节点是红色
                Node<K> uncle = gparent.right;
                if (uncle.color == RED) {
                    // Case 1条件：叔父节点是红色
                    // 叔父节点是红色，同时说明叔父节点不是NIL节点，修改颜色
                    uncle.color = BLACK;
                    parent.color = BLACK;
                    gparent.color = RED;
                    // 祖父节点不会是NIL节点，更不会是null
                    node = gparent;
                } else {
                    // 叔父节点是黑色，叔父可能是一个NIL节点，也可能不是。但不影响处理
                    // Case 2条件：叔父是黑色，且当前节点是右孩子，先进行左旋，然后执行case 3
                    if (parent.right == node) {
                        // 左旋
                        leftRotate(parent);
                        // 左旋后，node和parent的关系颠倒了，调整node和parent引用
                        node = parent;
                        parent = node.parent;
                    }

                    // 如果当前节点是左孩子，直接进行右旋
                    // Case 3条件：叔叔是黑色，且当前节点是左孩子。
                    parent.color = BLACK;
                    gparent.color = RED;
                    rightRotate(gparent);
                }
            } else {    //若“z的父节点”是“z的祖父节点的右孩子”
                // Case 1条件：叔叔节点是红色
                Node<K> uncle = gparent.left;
                //if ((uncle!=null) && isRed(uncle)) {
                if (uncle.color == RED) {
                    uncle.color = BLACK;
                    parent.color = BLACK;
                    gparent.color = RED;
                    node = gparent;
                } else {
                    // Case 2条件：叔叔是黑色，且当前节点是左孩子
                    if (parent.left == node) {
                        rightRotate(parent);
                        node = parent;
                        parent = node.parent;
                    }

                    // Case 3条件：叔叔是黑色，且当前节点是右孩子。
                    parent.color = BLACK;
                    gparent.color = RED;
                    leftRotate(gparent);
                }
            }
        }

        // 将根节点设为黑色
        root.color = BLACK;
    }

    private void insert(Node<K> node) {
        int cmp;
        Node<K> y = NIL;
        Node<K> x = root;

        // 1. 将红黑树当作一颗二叉查找树，将节点添加到二叉查找树中。
        while (x != NIL) {
            y = x;
            cmp = node.key.compareTo(x.key);
            if (cmp < 0) {
                x = x.left;
            } else {
                x = x.right;
            }
        }

        node.parent = y;
        if (y != NIL) {
            cmp = node.key.compareTo(y.key);
            if (cmp < 0) {
                y.left = node;
            } else {
                y.right = node;
            }
        } else {
            root = node;
        }

        // 2. 设置节点的颜色为红色
        node.color = RED;

        // 3. 将它重新修正为一颗二叉查找树
        insertFixUp(node);
    }

    public void insert(K key) {
        Node<K> node = new Node<>(key, BLACK, NIL, NIL, NIL);
        insert(node);
    }

    private void removeFixUp(Node<K> node) {

        // 兄弟节点
        Node<K> sibling;
        Node<K> parent = node.parent;
        while (node.color == BLACK && node != root) {
            if (parent.left == node) {
                sibling = parent.right;
                if (sibling.color == RED) {
                    // Case 1: x的兄弟w是红色的
                    sibling.color = BLACK;
                    parent.color = RED;
                    leftRotate(parent);
                    sibling = parent.right;
                }

                if (sibling.left.color == BLACK && sibling.right.color == BLACK) {
                    // Case 2: x的兄弟w是黑色，且w的俩个孩子也都是黑色的
                    sibling.color = RED;
                    node = parent;
                } else {
                    if (sibling.right.color == BLACK) {
                        // Case 3: x的兄弟w是黑色的，并且w的左孩子是红色，右孩子为黑色。
                        sibling.left.color = BLACK;
                        sibling.color = RED;
                        rightRotate(sibling);
                        sibling = parent.right;
                    }
                    // Case 4: x的兄弟w是黑色的；并且w的右孩子是红色的，左孩子任意颜色。
                    sibling.color = parent.color;
                    parent.color = BLACK;
                    sibling.right.color = BLACK;
                    leftRotate(parent);
                    node = root;
                }
            } else {
                sibling = parent.left;
                if (sibling.color == RED) {
                    // Case 1: x的兄弟w是红色的
                    sibling.color = BLACK;
                    parent.color = RED;
                    rightRotate(parent);
                    sibling = parent.left;
                }

                if (sibling.left.color == BLACK && sibling.right.color == BLACK) {
                    // Case 2: x的兄弟w是黑色，且w的俩个孩子也都是黑色的
                    sibling.color = RED;
                    node = parent;
                } else {

                    if (sibling.left.color == BLACK) {
                        // Case 3: x的兄弟w是黑色的，并且w的左孩子是红色，右孩子为黑色。
                        sibling.right.color = BLACK;
                        sibling.color = RED;
                        leftRotate(sibling);
                        sibling = parent.left;
                    }

                    // Case 4: x的兄弟w是黑色的；并且w的右孩子是红色的，左孩子任意颜色。
                    sibling.color = parent.color;
                    parent.color = BLACK;
                    sibling.left.color = BLACK;
                    rightRotate(parent);
                    node = root;
                }
            }
        }

        node.color = BLACK;
    }

    /**
     * 将替换节点放到被删除节点在树中的位置，修改替换节点的父节点为被删除节点的父节点。
     * 但是他们的孩子节点关系没有在这个方法里调整
     * @param removedNode 被删除的节点
     * @param replaceNode 用来替换被删除节点的节点
     */
    private void transplant(Node<K> removedNode, Node<K> replaceNode) {

        if (removedNode.parent == NIL) {
            // 删除的是root节点，替换节点变为root节点
            // root节点的父节点是NIL节点，但NIL节点的孩子并不是root
            root = replaceNode;
        } else if ( removedNode == removedNode.parent.left) {
            // 被删除节点是父节点的左孩子
            removedNode.parent.left = replaceNode;
        } else {
            // 被删除节点是父节点的右孩子
            removedNode.parent.right = replaceNode;
        }
        // 替换节点的父节点调整为被删除节点的父节点
        replaceNode.parent = removedNode.parent;

    }

    private void remove(Node<K> node) {
        // 需要调整的节点
        Node<K> fixupNode;
        // 从树中移除了的节点颜色
        boolean removedColor;

        if (node.left == NIL) {
            // 只有右孩子，直接将右孩子移到被删除节点位置
            removedColor = node.color;
            fixupNode = node.right;
            transplant(node, node.right);
        } else if (node.right == NIL) {
            // 只有左孩子，直接将左孩子移到被删除节点位置
            removedColor = node.color;
            fixupNode = node.left;
            transplant(node, node.left);
        } else {
            // 被删除节点有左右两个孩子
            // 将被删除节点的直接后继移动到被删除节点的位置
            Node<K> replace = minimum(node.right);
            removedColor = replace.color;
            fixupNode = replace.right;

            // 调整replace、fixupnode的parent、孩子关系
            if (replace.parent == node) {
                fixupNode.parent = replace;
            } else {
                transplant(replace, replace.right);
                replace.right = node.right;
                replace.right.parent = replace;
            }

            transplant(node, replace);
            replace.left = node.left;
            replace.left.parent = replace;
            replace.color = node.color;
        }

        // 如果被删除节点路径上减少了黑节点，表示降低了黑高，需要进行调整
        if (removedColor == BLACK) {
            removeFixUp(fixupNode);
        }
    }

    public void remove(K key) {
        Node<K> node = search(root, key);
        if (node != NIL) {
            remove(node);
        }
    }

    private void destroy(Node<K> tree) {
        if (tree == NIL)
            return ;

        if (tree.left != NIL)
            destroy(tree.left);
        if (tree.right != NIL)
            destroy(tree.right);

        tree = null;
    }

    public void clear() {
        destroy(root);
        root = null;
    }

    /**
     * 打印"红黑树"
     *
     * key        -- 节点的键值
     * direction  --  0，表示该节点是根节点;
     *               -1，表示该节点是它的父结点的左孩子;
     *                1，表示该节点是它的父结点的右孩子。
     */
    private void print(Node<K> tree, K key, int direction) {

        if (tree != NIL) {

            if (direction == 0) {
                // tree是根节点
                System.out.printf("%2d(B) is root\n", tree.key);
            } else {
                // tree是分支节点
                System.out.printf("%2d(%s) is %2d's %6s child\n", tree.key, tree.color == RED ? "R" : "B",
                        key, direction == 1 ? "right" : "left");
            }
            print(tree.left, tree.key, -1);
            print(tree.right,tree.key,  1);
        }
    }

    public void print() {
        if (root != NIL) {
            print(root, root.key, 0);
        }
    }

    public static class Node<K extends Comparable<K>> {
        private boolean color;        // 颜色
        private K key;                // 关键字(键值)
        private Node<K> left;    // 左孩子
        private Node<K> right;    // 右孩子
        private Node<K> parent;    // 父结点

        public Node(K key, boolean color, Node<K> parent, Node<K> left, Node<K> right) {
            this.key = key;
            this.color = color;
            this.parent = parent;
            this.left = left;
            this.right = right;
        }

        public String toString() {
            return "" + key + (this.color == RED ? "(R)" : "(B)");
        }
    }

}
